{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# 导包&GPU",
   "id": "ab1d7004b72ab354"
  },
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# Continue with regular imports\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from torch import nn\n",
    "\n",
    "import torchvision\n",
    "\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\""
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# 准备数据集",
   "id": "5c3bb4c3e4af932a"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "import os\n",
    "import zipfile\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "import requests\n",
    "\n",
    "def download_data(source: str, \n",
    "                  destination: str,\n",
    "                  remove_source: bool = True) -> Path:\n",
    "    \"\"\"Downloads a zipped dataset from source and unzips to destination.\n",
    "\n",
    "    Args:\n",
    "        source (str): A link to a zipped file containing data.\n",
    "        destination (str): A target directory to unzip data to.\n",
    "        remove_source (bool): Whether to remove the source after downloading and extracting.\n",
    "    \n",
    "    Returns:\n",
    "        pathlib.Path to downloaded data.\n",
    "    \n",
    "    Example usage:\n",
    "        download_data(source=\"https://github.com/mrdbourke/pytorch-deep-learning/raw/main/data/pizza_steak_sushi.zip\",\n",
    "                      destination=\"pizza_steak_sushi\")\n",
    "    \"\"\"\n",
    "    # Setup path to data folder\n",
    "    data_path = Path(\"dataset/\")\n",
    "    image_path = data_path / destination\n",
    "\n",
    "    # If the image folder doesn't exist, download it and prepare it... \n",
    "    if image_path.is_dir():\n",
    "        print(f\"[INFO] {image_path} directory exists, skipping download.\")\n",
    "    else:\n",
    "        print(f\"[INFO] Did not find {image_path} directory, creating one...\")\n",
    "        image_path.mkdir(parents=True, exist_ok=True)\n",
    "        \n",
    "        # Download pizza, steak, sushi data\n",
    "        target_file = Path(source).name\n",
    "        with open(data_path / target_file, \"wb\") as f:\n",
    "            request = requests.get(source)\n",
    "            print(f\"[INFO] Downloading {target_file} from {source}...\")\n",
    "            f.write(request.content)\n",
    "\n",
    "        # Unzip pizza, steak, sushi data\n",
    "        with zipfile.ZipFile(data_path / target_file, \"r\") as zip_ref:\n",
    "            print(f\"[INFO] Unzipping {target_file} data...\") \n",
    "            zip_ref.extractall(image_path)\n",
    "\n",
    "        # Remove .zip file\n",
    "        if remove_source:\n",
    "            os.remove(data_path / target_file)\n",
    "    \n",
    "    return image_path\n",
    "\n",
    "image_path = download_data(source=\"https://github.com/mrdbourke/pytorch-deep-learning/raw/main/data/pizza_steak_sushi.zip\",\n",
    "                           destination=\"pizza_steak_sushi\")\n",
    "image_path"
   ],
   "id": "7229aa57f7496b5",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# 设置随机种子",
   "id": "82613d22f06b9f51"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# Set seeds\n",
    "def set_seeds(seed: int=42):\n",
    "    \"\"\"Sets random sets for torch operations.\n",
    "    Args:\n",
    "        seed (int, optional): Random seed to set. Defaults to 42.\n",
    "    \"\"\"\n",
    "    # Set the seed for general torch operations\n",
    "    torch.manual_seed(seed)\n",
    "    # Set the seed for CUDA torch operations (ones that happen on the GPU)\n",
    "    torch.cuda.manual_seed(seed)"
   ],
   "id": "e6294b9d667ecb39",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# Datasets & DataLoaders",
   "id": "372016de64f6f491"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "import data_setup\n",
    "\n",
    "# Setup directories\n",
    "train_dir = image_path / \"train\"\n",
    "test_dir = image_path / \"test\"\n",
    "\n",
    "# Setup ImageNet normalization levels (turns all images into similar distribution as ImageNet)\n",
    "normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                                 std=[0.229, 0.224, 0.225])\n",
    "\n",
    "# Create transform pipeline manually\n",
    "manual_transforms = transforms.Compose([\n",
    "    transforms.Resize((224, 224)),\n",
    "    transforms.ToTensor(),\n",
    "    normalize\n",
    "])           \n",
    "print(f\"Manually created transforms: {manual_transforms}\")\n",
    "\n",
    "# Create data loaders\n",
    "train_dataloader, test_dataloader, class_names = data_setup.create_dataloaders(\n",
    "    train_dir=train_dir,\n",
    "    test_dir=test_dir,\n",
    "    transform=manual_transforms, # use manually created transforms\n",
    "    batch_size=32\n",
    ")\n",
    "\n",
    "train_dataloader, test_dataloader, class_names"
   ],
   "id": "758c3ff2a42faf88",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# 准备（预训练）模型",
   "id": "b3c36ec28988f922"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# model = torchvision.models.efficientnet_b0(pretrained=True).to(device) # OLD \n",
    "\n",
    "# Download the pretrained weights for EfficientNet_B0\n",
    "weights = torchvision.models.EfficientNet_B0_Weights.DEFAULT # NEW in torchvision 0.13, \"DEFAULT\" means \"best weights available\"\n",
    "\n",
    "# Setup the model with the pretrained weights and send it to the target device\n",
    "model = torchvision.models.efficientnet_b0(weights=weights).to(device)\n",
    "\n",
    "# View the output of the model\n",
    "model"
   ],
   "id": "d0ae8ff143672d3e",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## 设置requires_grad为False以冻结预训练模型基础层参数",
   "id": "11e80853e7f2b7b"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "for param in model.features.parameters():\n",
    "    param.requires_grad = False\n",
    "    \n",
    "# Since we're creating a new layer with random weights (torch.nn.Linear), \n",
    "# let's set the seeds\n",
    "set_seeds() \n",
    "# Update the classifier head to suit our problem\n",
    "model.classifier = torch.nn.Sequential(\n",
    "    nn.Dropout(p=0.2, inplace=True),\n",
    "    nn.Linear(in_features=1280, \n",
    "              out_features=len(class_names),\n",
    "              bias=True).to(device))"
   ],
   "id": "2cf445b47b3c7552",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## summary",
   "id": "bed393171690cf68"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "from torchinfo import summary\n",
    "\n",
    "# # Get a summary of the model (uncomment for full output)\n",
    "summary(model, \n",
    "        input_size=(32, 3, 224, 224), # make sure this is \"input_size\", not \"input_shape\" (batch_size, color_channels, height, width)\n",
    "        verbose=0,\n",
    "        col_names=[\"input_size\", \"output_size\", \"num_params\", \"trainable\"],\n",
    "        col_width=20,\n",
    "        row_settings=[\"var_names\"]\n",
    ")"
   ],
   "id": "4e561ed86a44bfcd",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## loss & optimizer",
   "id": "5aa55a16a0f1e36b"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# Define loss and optimizer\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)"
   ],
   "id": "10862c4565365a8a",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## SummaryWriter",
   "id": "b699303e2f5f53c"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "from torch.utils.tensorboard import SummaryWriter\n",
    "\n",
    "# Create a writer with all default settings\n",
    "writer = SummaryWriter()"
   ],
   "id": "641b1070a5b7639f",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "from typing import Dict, List\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "from engine import train_step, test_step\n",
    "\n",
    "def train(model: torch.nn.Module, \n",
    "          train_dataloader: torch.utils.data.DataLoader, \n",
    "          test_dataloader: torch.utils.data.DataLoader, \n",
    "          optimizer: torch.optim.Optimizer,\n",
    "          loss_fn: torch.nn.Module,\n",
    "          epochs: int,\n",
    "          device: torch.device) -> Dict[str, List]:\n",
    "    \"\"\"Trains and tests a PyTorch model.\n",
    "    Args:\n",
    "      model: A PyTorch model to be trained and tested.\n",
    "      train_dataloader: A DataLoader instance for the model to be trained on.\n",
    "      test_dataloader: A DataLoader instance for the model to be tested on.\n",
    "      optimizer: A PyTorch optimizer to help minimize the loss function.\n",
    "      loss_fn: A PyTorch loss function to calculate loss on both datasets.\n",
    "      epochs: An integer indicating how many epochs to train for.\n",
    "      device: A target device to compute on (e.g. \"cuda\" or \"cpu\").\n",
    "      \n",
    "    Returns:\n",
    "      A dictionary of training and testing loss as well as training and\n",
    "      testing accuracy metrics. Each metric has a value in a list for \n",
    "      each epoch.\n",
    "      In the form: {train_loss: [...],\n",
    "                train_acc: [...],\n",
    "                test_loss: [...],\n",
    "                test_acc: [...]} \n",
    "      For example if training for epochs=2: \n",
    "              {train_loss: [2.0616, 1.0537],\n",
    "                train_acc: [0.3945, 0.3945],\n",
    "                test_loss: [1.2641, 1.5706],\n",
    "                test_acc: [0.3400, 0.2973]} \n",
    "    \"\"\"\n",
    "    # Create empty results dictionary\n",
    "    results = {\"train_loss\": [],\n",
    "               \"train_acc\": [],\n",
    "               \"test_loss\": [],\n",
    "               \"test_acc\": []\n",
    "    }\n",
    "\n",
    "    # Loop through training and testing steps for a number of epochs\n",
    "    for epoch in tqdm(range(epochs)):\n",
    "        train_loss, train_acc = train_step(model=model,\n",
    "                                           dataloader=train_dataloader,\n",
    "                                           loss_fn=loss_fn,\n",
    "                                           optimizer=optimizer,\n",
    "                                           device=device)\n",
    "        test_loss, test_acc = test_step(model=model,\n",
    "                                        dataloader=test_dataloader,\n",
    "                                        loss_fn=loss_fn,\n",
    "                                        device=device)\n",
    "\n",
    "        # Print out what's happening\n",
    "        print(\n",
    "          f\"Epoch: {epoch+1} | \"\n",
    "          f\"train_loss: {train_loss:.4f} | \"\n",
    "          f\"train_acc: {train_acc:.4f} | \"\n",
    "          f\"test_loss: {test_loss:.4f} | \"\n",
    "          f\"test_acc: {test_acc:.4f}\"\n",
    "        )\n",
    "\n",
    "        # Update results dictionary\n",
    "        results[\"train_loss\"].append(train_loss)\n",
    "        results[\"train_acc\"].append(train_acc)\n",
    "        results[\"test_loss\"].append(test_loss)\n",
    "        results[\"test_acc\"].append(test_acc)\n",
    "\n",
    "        ### New: Experiment tracking ###\n",
    "        # Add loss results to SummaryWriter\n",
    "        writer.add_scalars(main_tag=\"Loss\", \n",
    "                           tag_scalar_dict={\"train_loss\": train_loss,\n",
    "                                            \"test_loss\": test_loss},\n",
    "                           global_step=epoch)\n",
    "\n",
    "        # Add accuracy results to SummaryWriter\n",
    "        writer.add_scalars(main_tag=\"Accuracy\", \n",
    "                           tag_scalar_dict={\"train_acc\": train_acc,\n",
    "                                            \"test_acc\": test_acc}, \n",
    "                           global_step=epoch)\n",
    "        \n",
    "        # Track the PyTorch model architecture\n",
    "        writer.add_graph(model=model, \n",
    "                         # Pass in an example input\n",
    "                         input_to_model=torch.randn(32, 3, 224, 224).to(device))\n",
    "    \n",
    "    # Close the writer\n",
    "    writer.close()\n",
    "    \n",
    "    ### End new ###\n",
    "\n",
    "    # Return the filled results at the end of the epochs\n",
    "    return results"
   ],
   "id": "feed955d92dcf90c",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# 训练",
   "id": "561b4d70be73da7f"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "set_seeds()\n",
    "results = train(model=model,\n",
    "                train_dataloader=train_dataloader,\n",
    "                test_dataloader=test_dataloader,\n",
    "                optimizer=optimizer,\n",
    "                loss_fn=loss_fn,\n",
    "                epochs=5,\n",
    "                device=device)"
   ],
   "id": "1a7d88a1d7eecbc6",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "> 产生了一个`runs/`目录存储模型训练结果`runs/CURRENT_DATETIME_HOSTNAME`",
   "id": "afa9e06ecdb0a69b"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# 在TensorBoard中可视化查看模型结果",
   "id": "67246c9ddceda27"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "%load_ext tensorboard # line magic to load TensorBoard\n",
    "%tensorboard --logdir runs # run TensorBoard session with the \"runs/\" directory\n",
    "\n",
    "# pycharm中打开命令行输入`tensorboard --logdir=<directory_name>`"
   ],
   "id": "ef73fab34cb138bc",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# `create_writer()`",
   "id": "6a435a2c76e21e89"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "> runs/YYYY-MM-DD/experiment_name/model_name/extra",
   "id": "f0227ad47acb4ad2"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "def create_writer(experiment_name: str, \n",
    "                  model_name: str, \n",
    "                  extra: str=None) -> torch.utils.tensorboard.writer.SummaryWriter():\n",
    "    \"\"\"Creates a torch.utils.tensorboard.writer.SummaryWriter() instance saving to a specific log_dir.\n",
    "\n",
    "    log_dir is a combination of runs/timestamp/experiment_name/model_name/extra.\n",
    "\n",
    "    Where timestamp is the current date in YYYY-MM-DD format.\n",
    "\n",
    "    Args:\n",
    "        experiment_name (str): Name of experiment.\n",
    "        model_name (str): Name of model.\n",
    "        extra (str, optional): Anything extra to add to the directory. Defaults to None.\n",
    "\n",
    "    Returns:\n",
    "        torch.utils.tensorboard.writer.SummaryWriter(): Instance of a writer saving to log_dir.\n",
    "\n",
    "    Example usage:\n",
    "        # Create a writer saving to \"runs/2022-06-04/data_10_percent/effnetb2/5_epochs/\"\n",
    "        writer = create_writer(experiment_name=\"data_10_percent\",\n",
    "                               model_name=\"effnetb2\",\n",
    "                               extra=\"5_epochs\")\n",
    "        # The above is the same as:\n",
    "        writer = SummaryWriter(log_dir=\"runs/2022-06-04/data_10_percent/effnetb2/5_epochs/\")\n",
    "    \"\"\"\n",
    "    from datetime import datetime\n",
    "    import os\n",
    "\n",
    "    # Get timestamp of current date (all experiments on certain day live in same folder)\n",
    "    timestamp = datetime.now().strftime(\"%Y-%m-%d\") # returns current date in YYYY-MM-DD format\n",
    "\n",
    "    if extra:\n",
    "        # Create log directory path\n",
    "        log_dir = os.path.join(\"runs\", timestamp, experiment_name, model_name, extra)\n",
    "    else:\n",
    "        log_dir = os.path.join(\"runs\", timestamp, experiment_name, model_name)\n",
    "        \n",
    "    print(f\"[INFO] Created SummaryWriter, saving to: {log_dir}...\")\n",
    "    return SummaryWriter(log_dir=log_dir)"
   ],
   "id": "197a01112c842246",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# Create an example writer\n",
    "example_writer = create_writer(experiment_name=\"data_10_percent\",\n",
    "                               model_name=\"effnetb0\",\n",
    "                               extra=\"5_epochs\")"
   ],
   "id": "77867381dfcb384f",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## 对不同模型的每次实验设置不同的`writer`实例",
   "id": "6ba0feff59329bce"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "from typing import Dict, List\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "# Add writer parameter to train()\n",
    "def train(model: torch.nn.Module, \n",
    "          train_dataloader: torch.utils.data.DataLoader, \n",
    "          test_dataloader: torch.utils.data.DataLoader, \n",
    "          optimizer: torch.optim.Optimizer,\n",
    "          loss_fn: torch.nn.Module,\n",
    "          epochs: int,\n",
    "          device: torch.device, \n",
    "          writer: torch.utils.tensorboard.writer.SummaryWriter # new \n",
    "          ) -> Dict[str, List]:\n",
    "\n",
    "    # Create empty results dictionary\n",
    "    results = {\"train_loss\": [],\n",
    "               \"train_acc\": [],\n",
    "               \"test_loss\": [],\n",
    "               \"test_acc\": []\n",
    "    }\n",
    "\n",
    "    # Loop through training and testing steps for a number of epochs\n",
    "    for epoch in tqdm(range(epochs)):\n",
    "        train_loss, train_acc = train_step(model=model,\n",
    "                                          dataloader=train_dataloader,\n",
    "                                          loss_fn=loss_fn,\n",
    "                                          optimizer=optimizer,\n",
    "                                          device=device)\n",
    "        test_loss, test_acc = test_step(model=model,\n",
    "          dataloader=test_dataloader,\n",
    "          loss_fn=loss_fn,\n",
    "          device=device)\n",
    "\n",
    "        # Print out what's happening\n",
    "        print(\n",
    "          f\"Epoch: {epoch+1} | \"\n",
    "          f\"train_loss: {train_loss:.4f} | \"\n",
    "          f\"train_acc: {train_acc:.4f} | \"\n",
    "          f\"test_loss: {test_loss:.4f} | \"\n",
    "          f\"test_acc: {test_acc:.4f}\"\n",
    "        )\n",
    "\n",
    "        # Update results dictionary\n",
    "        results[\"train_loss\"].append(train_loss)\n",
    "        results[\"train_acc\"].append(train_acc)\n",
    "        results[\"test_loss\"].append(test_loss)\n",
    "        results[\"test_acc\"].append(test_acc)\n",
    "\n",
    "\n",
    "        ### New: Use the writer parameter to track experiments ###\n",
    "        # See if there's a writer, if so, log to it\n",
    "        if writer:\n",
    "            # Add results to SummaryWriter\n",
    "            writer.add_scalars(main_tag=\"Loss\", \n",
    "                               tag_scalar_dict={\"train_loss\": train_loss,\n",
    "                                                \"test_loss\": test_loss},\n",
    "                               global_step=epoch)\n",
    "            writer.add_scalars(main_tag=\"Accuracy\", \n",
    "                               tag_scalar_dict={\"train_acc\": train_acc,\n",
    "                                                \"test_acc\": test_acc}, \n",
    "                               global_step=epoch)\n",
    "\n",
    "            # Close the writer\n",
    "            writer.close()\n",
    "        else:\n",
    "            pass\n",
    "    ### End new ###\n",
    "\n",
    "    # Return the filled results at the end of the epochs\n",
    "    return results"
   ],
   "id": "ef2b327a141e1f18",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# 模型比较",
   "id": "834de2dc7f86d38a"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "> - 不同的数据量(10% vs. 20%)\n",
    "> - 不同的模型(`torchvision.models.efficientnet_b0` vs. `torchvision.models.efficientnet_b2`)\n",
    "> - 不同的训练次数(5 epochs vs. 10 epochs)\n",
    "\n",
    "**尤其注意，在实验中使用10%-20%的少量数据集，如果有效再全量训练**"
   ],
   "id": "96c3790fed1e7864"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# Download 10 percent and 20 percent training data (if necessary)\n",
    "data_10_percent_path = download_data(source=\"https://github.com/mrdbourke/pytorch-deep-learning/raw/main/data/pizza_steak_sushi.zip\",\n",
    "                                     destination=\"pizza_steak_sushi\")\n",
    "\n",
    "data_20_percent_path = download_data(source=\"https://github.com/mrdbourke/pytorch-deep-learning/raw/main/data/pizza_steak_sushi_20_percent.zip\",\n",
    "                                     destination=\"pizza_steak_sushi_20_percent\")"
   ],
   "id": "3c3f0989b3525268",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# Setup training directory paths\n",
    "train_dir_10_percent = data_10_percent_path / \"train\"\n",
    "train_dir_20_percent = data_20_percent_path / \"train\"\n",
    "\n",
    "# Setup testing directory paths (note: use the same test dataset for both to compare the results)\n",
    "test_dir = data_10_percent_path / \"test\"\n",
    "\n",
    "# Check the directories\n",
    "print(f\"Training directory 10%: {train_dir_10_percent}\")\n",
    "print(f\"Training directory 20%: {train_dir_20_percent}\")\n",
    "print(f\"Testing directory: {test_dir}\")"
   ],
   "id": "bc7280e087a7fc07",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## transforms",
   "id": "25affb568b4cb1f9"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "from torchvision import transforms\n",
    "\n",
    "# Create a transform to normalize data distribution to be inline with ImageNet\n",
    "normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], # values per colour channel [red, green, blue]\n",
    "                                 std=[0.229, 0.224, 0.225]) # values per colour channel [red, green, blue]\n",
    "\n",
    "# Compose transforms into a pipeline\n",
    "simple_transform = transforms.Compose([\n",
    "    transforms.Resize((224, 224)), # 1. Resize the images\n",
    "    transforms.ToTensor(), # 2. Turn the images into tensors with values between 0 & 1\n",
    "    normalize # 3. Normalize the images so their distributions match the ImageNet dataset \n",
    "])"
   ],
   "id": "a30d33213017031d",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## dataset & dataloader",
   "id": "cfe5af4a871d7aad"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "BATCH_SIZE = 32\n",
    "\n",
    "# Create 10% training and test DataLoaders\n",
    "train_dataloader_10_percent, test_dataloader, class_names = data_setup.create_dataloaders(train_dir=train_dir_10_percent,\n",
    "    test_dir=test_dir, \n",
    "    transform=simple_transform,\n",
    "    batch_size=BATCH_SIZE\n",
    ")\n",
    "\n",
    "# Create 20% training and test data DataLoders\n",
    "train_dataloader_20_percent, test_dataloader, class_names = data_setup.create_dataloaders(train_dir=train_dir_20_percent,\n",
    "    test_dir=test_dir,\n",
    "    transform=simple_transform,\n",
    "    batch_size=BATCH_SIZE\n",
    ")\n",
    "\n",
    "# Find the number of samples/batches per dataloader (using the same test_dataloader for both experiments)\n",
    "print(f\"Number of batches of size {BATCH_SIZE} in 10 percent training data: {len(train_dataloader_10_percent)}\")\n",
    "print(f\"Number of batches of size {BATCH_SIZE} in 20 percent training data: {len(train_dataloader_20_percent)}\")\n",
    "print(f\"Number of batches of size {BATCH_SIZE} in testing data: {len(train_dataloader_10_percent)} (all experiments will use the same test set)\")\n",
    "print(f\"Number of classes: {len(class_names)}, class names: {class_names}\")"
   ],
   "id": "975a8ceadeaf1d20",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## 获取`EffNetB2`",
   "id": "f421c71af15a2aeb"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "import torchvision\n",
    "from torchinfo import summary\n",
    "\n",
    "# 1. Create an instance of EffNetB2 with pretrained weights\n",
    "effnetb2_weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT # \"DEFAULT\" means best available weights\n",
    "effnetb2 = torchvision.models.efficientnet_b2(weights=effnetb2_weights)\n",
    "\n",
    "# # 2. Get a summary of standard EffNetB2 from torchvision.models (uncomment for full output)\n",
    "summary(model=effnetb2, \n",
    "        input_size=(32, 3, 224, 224), # make sure this is \"input_size\", not \"input_shape\"\n",
    "        # col_names=[\"input_size\"], # uncomment for smaller output\n",
    "        col_names=[\"input_size\", \"output_size\", \"num_params\", \"trainable\"],\n",
    "        col_width=20,\n",
    "        row_settings=[\"var_names\"]\n",
    ") \n",
    "\n",
    "# 3. Get the number of in_features of the EfficientNetB2 classifier layer\n",
    "print(f\"Number of in_features to final layer of EfficientNetB2: {len(effnetb2.classifier.state_dict()['1.weight'][0])}\")"
   ],
   "id": "a9b782f4de29f53d",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "import torchvision\n",
    "from torch import nn\n",
    "\n",
    "# Get num out features (one for each class pizza, steak, sushi)\n",
    "OUT_FEATURES = len(class_names)\n",
    "\n",
    "# Create an EffNetB0 feature extractor\n",
    "def create_effnetb0():\n",
    "    # 1. Get the base mdoel with pretrained weights and send to target device\n",
    "    weights = torchvision.models.EfficientNet_B0_Weights.DEFAULT\n",
    "    model = torchvision.models.efficientnet_b0(weights=weights).to(device)\n",
    "\n",
    "    # 2. Freeze the base model layers\n",
    "    for param in model.features.parameters():\n",
    "        param.requires_grad = False\n",
    "\n",
    "    # 3. Set the seeds\n",
    "    set_seeds()\n",
    "\n",
    "    # 4. Change the classifier head\n",
    "    model.classifier = nn.Sequential(\n",
    "        nn.Dropout(p=0.2),\n",
    "        nn.Linear(in_features=1280, out_features=OUT_FEATURES)\n",
    "    ).to(device)\n",
    "\n",
    "    # 5. Give the model a name\n",
    "    model.name = \"effnetb0\"\n",
    "    print(f\"[INFO] Created new {model.name} model.\")\n",
    "    return model\n",
    "\n",
    "# Create an EffNetB2 feature extractor\n",
    "def create_effnetb2():\n",
    "    # 1. Get the base model with pretrained weights and send to target device\n",
    "    weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT\n",
    "    model = torchvision.models.efficientnet_b2(weights=weights).to(device)\n",
    "\n",
    "    # 2. Freeze the base model layers\n",
    "    for param in model.features.parameters():\n",
    "        param.requires_grad = False\n",
    "\n",
    "    # 3. Set the seeds\n",
    "    set_seeds()\n",
    "\n",
    "    # 4. Change the classifier head\n",
    "    model.classifier = nn.Sequential(\n",
    "        nn.Dropout(p=0.3),\n",
    "        nn.Linear(in_features=1408, out_features=OUT_FEATURES)\n",
    "    ).to(device)\n",
    "\n",
    "    # 5. Give the model a name\n",
    "    model.name = \"effnetb2\"\n",
    "    print(f\"[INFO] Created new {model.name} model.\")\n",
    "    return model"
   ],
   "id": "fb390b489fa56575",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "effnetb0 = create_effnetb0() \n",
    "\n",
    "# Get an output summary of the layers in our EffNetB0 feature extractor model (uncomment to view full output)\n",
    "# summary(model=effnetb0, \n",
    "#         input_size=(32, 3, 224, 224), # make sure this is \"input_size\", not \"input_shape\"\n",
    "#         # col_names=[\"input_size\"], # uncomment for smaller output\n",
    "#         col_names=[\"input_size\", \"output_size\", \"num_params\", \"trainable\"],\n",
    "#         col_width=20,\n",
    "#         row_settings=[\"var_names\"]\n",
    "# ) \n",
    "\n",
    "effnetb2 = create_effnetb2()\n",
    "\n",
    "# Get an output summary of the layers in our EffNetB2 feature extractor model (uncomment to view full output)\n",
    "# summary(model=effnetb2, \n",
    "#         input_size=(32, 3, 224, 224), # make sure this is \"input_size\", not \"input_shape\"\n",
    "#         # col_names=[\"input_size\"], # uncomment for smaller output\n",
    "#         col_names=[\"input_size\", \"output_size\", \"num_params\", \"trainable\"],\n",
    "#         col_width=20,\n",
    "#         row_settings=[\"var_names\"]\n",
    "# ) "
   ],
   "id": "9d06139a3170278a",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## 对比训练",
   "id": "48dce781daad3bd6"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# 1. Create epochs list\n",
    "num_epochs = [5, 10]\n",
    "\n",
    "# 2. Create models list (need to create a new model for each experiment)\n",
    "models = [\"effnetb0\", \"effnetb2\"]\n",
    "\n",
    "# 3. Create dataloaders dictionary for various dataloaders\n",
    "train_dataloaders = {\"data_10_percent\": train_dataloader_10_percent,\n",
    "                     \"data_20_percent\": train_dataloader_20_percent}"
   ],
   "id": "48479c1cbe6b8fbc",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "%%time\n",
    "from utils import save_model\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "print(device)\n",
    "# 1. Set the random seeds\n",
    "set_seeds(seed=42)\n",
    "\n",
    "# 2. Keep track of experiment numbers\n",
    "experiment_number = 0\n",
    "\n",
    "# 3. Loop through each DataLoader\n",
    "for dataloader_name, train_dataloader in train_dataloaders.items():\n",
    "\n",
    "    # 4. Loop through each number of epochs\n",
    "    for epochs in num_epochs: \n",
    "\n",
    "        # 5. Loop through each model name and create a new model based on the name\n",
    "        for model_name in models:\n",
    "\n",
    "            # 6. Create information print outs\n",
    "            experiment_number += 1\n",
    "            print(f\"[INFO] Experiment number: {experiment_number}\")\n",
    "            print(f\"[INFO] Model: {model_name}\")\n",
    "            print(f\"[INFO] DataLoader: {dataloader_name}\")\n",
    "            print(f\"[INFO] Number of epochs: {epochs}\")  \n",
    "\n",
    "            # 7. Select the model\n",
    "            if model_name == \"effnetb0\":\n",
    "                model = create_effnetb0() # creates a new model each time (important because we want each experiment to start from scratch)\n",
    "            else:\n",
    "                model = create_effnetb2() # creates a new model each time (important because we want each experiment to start from scratch)\n",
    "            \n",
    "            # 8. Create a new loss and optimizer for every model\n",
    "            loss_fn = nn.CrossEntropyLoss()\n",
    "            optimizer = torch.optim.Adam(params=model.parameters(), lr=0.001)\n",
    "\n",
    "            # 9. Train target model with target dataloaders and track experiments\n",
    "            train(model=model,\n",
    "                  train_dataloader=train_dataloader,\n",
    "                  test_dataloader=test_dataloader, \n",
    "                  optimizer=optimizer,\n",
    "                  loss_fn=loss_fn,\n",
    "                  epochs=epochs,\n",
    "                  device=device,\n",
    "                  writer=create_writer(experiment_name=dataloader_name,\n",
    "                                       model_name=model_name,\n",
    "                                       extra=f\"{epochs}_epochs\"))\n",
    "            \n",
    "            # 10. Save the model to file so we can get back the best model\n",
    "            save_filepath = f\"07_{model_name}_{dataloader_name}_{epochs}_epochs.pth\"\n",
    "            save_model(model=model,\n",
    "                       target_dir=\"models\",\n",
    "                       model_name=save_filepath)\n",
    "            print(\"-\"*50 + \"\\n\")"
   ],
   "id": "dce765c543f669b4"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# 加载最优模型进行预测",
   "id": "ba64a20f588b981c"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "# Setup the best model filepath\n",
    "best_model_path = \"models/07_effnetb2_data_20_percent_10_epochs.pth\"\n",
    "\n",
    "# Instantiate a new instance of EffNetB2 (to load the saved state_dict() to)\n",
    "best_model = create_effnetb2()\n",
    "\n",
    "# Load the saved best model state_dict()\n",
    "best_model.load_state_dict(torch.load(best_model_path))"
   ],
   "id": "56f8e9f9dffdf93"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "# Check the model file size\n",
    "from pathlib import Path\n",
    "\n",
    "# Get the model size in bytes then convert to megabytes\n",
    "effnetb2_model_size = Path(best_model_path).stat().st_size // (1024*1024)\n",
    "print(f\"EfficientNetB2 feature extractor model size: {effnetb2_model_size} MB\")"
   ],
   "id": "30507135b3f30af1"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
